{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Staggered grids\n",
    "\n",
    "[![Google Collab Book](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tum-pbs/PhiFlow/blob/develop/docs/Staggered_Grids.ipynb)\n",
    "\n",
    "Staggered grids are a key component of the marker-and-cell (MAC) method \\[Harlow and Welch 1965\\].\n",
    "They sample the velocity components not at the cell centers but in staggered form at the corresponding face centers.\n",
    "Their main advantage is that the divergence of a cell can be computed exactly.\n",
    "\n",
    "Φ<sub>Flow</sub> only stores valid velocity values in memory.\n",
    "This may require non-uniform tensors for the values since the numbers of horizontal and vertical faces are generally not equal.\n",
    "Depending on the boundary conditions, the outer-most values may also be redundant and, thus, not stored.\n",
    "\n",
    "![image](./figures/Staggered.png)\n",
    "\n",
    "Φ<sub>Flow</sub> represents staggered grids as instances of [`StaggeredGrid`](phi/field/#phi.field.StaggeredGrid).\n",
    "They have the same properties as `CenteredGrid` but the `values` field may reference a\n",
    "[non-uniform tensor](https://tum-pbs.github.io/PhiFlow/Math.html#non-uniform-tensors)\n",
    "to reflect the varying number of x, y and z sample points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[92m(xˢ=\u001b[94m(x=11, y=10)\u001b[92m, yˢ=\u001b[94m(x=10, y=11)\u001b[92m, vectorᶜ=x,y)\u001b[0m \u001b[94mconst 0.0\u001b[0m"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# !pip install --quiet phiflow\n",
    "from phi.flow import *\n",
    "\n",
    "grid = StaggeredGrid(0, extrapolation.BOUNDARY, x=10, y=10)\n",
    "grid.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Here, each component of the values tensor has one more sample point in the direction it is facing.\n",
    "If the extrapolation was `extrapolation.ZERO`, it would be one less (see above image)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Creating Staggered Grids\n",
    "\n",
    "The [`StaggeredGrid` constructor](phi/field/#phi.field.StaggeredGrid) supports two modes:\n",
    "\n",
    "* **Direct construction** `StaggeredGrid(values: Tensor, extrapolation, bounds)`.\n",
    "  All required fields are passed as arguments and stored as-is.\n",
    "  The `values` tensor must have the correct shape considering the extrapolation.\n",
    "* **Construction by resampling** `StaggeredGrid(values: Any, extrapolation, bounds, resolution, **resolution)`.\n",
    "  When specifying the resolution as a `Shape` or via keyword arguments, non-Tensor values can be passed for `values`,\n",
    "  such as geometries, other fields, constants or functions (see the [documentation](phi/field/#phi.field.StaggeredGrid)).\n",
    "\n",
    "\n",
    "Examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "domain = dict(x=10, y=10, bounds=Box(x=1, y=1), extrapolation=extrapolation.ZERO)\n",
    "\n",
    "grid = StaggeredGrid((1, -1), **domain)  # from constant vector\n",
    "grid = StaggeredGrid(Noise(), **domain)  # sample analytic field\n",
    "grid = StaggeredGrid(grid, **domain)  # resample existing field\n",
    "grid = StaggeredGrid(lambda x: math.exp(-x), **domain)  # function value(location)\n",
    "grid = resample(Sphere(x=0, y=0, radius=1), StaggeredGrid(0, **domain))  # no anti-aliasing\n",
    "grid = resample(Sphere(x=0, y=0, radius=1), StaggeredGrid(0, **domain), soft=True)  # with anti-aliasing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To construct a `StaggeredGrid` from NumPy arrays (or TensorFlow/PyTorch/Jax tensors), the tensors first need to be converted to Φ<sub>Flow</sub> tensors using `tensor()` or `wrap()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StaggeredGrid[(xˢ=32, yˢ=32, vectorᶜ=2), size=\u001b[94m(x=32, y=32)\u001b[0m, extrapolation=\u001b[94m0\u001b[0m]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vx = tensor(np.zeros([33, 32]), spatial('x,y'))\n",
    "vy = tensor(np.zeros([32, 33]), spatial('x,y'))\n",
    "StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.BOUNDARY)\n",
    "\n",
    "vx = tensor(np.zeros([32, 32]), spatial('x,y'))\n",
    "vy = tensor(np.zeros([32, 32]), spatial('x,y'))\n",
    "StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), extrapolation.PERIODIC)\n",
    "\n",
    "vx = tensor(np.zeros([31, 32]), spatial('x,y'))\n",
    "vy = tensor(np.zeros([32, 31]), spatial('x,y'))\n",
    "StaggeredGrid(math.stack([vx, vy], dual(vector='x,y')), 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Staggered grids can also be created from other fields using `field.at()` or `@` by passing an existing `StaggeredGrid`.\n",
    "\n",
    "Some field functions also return `StaggeredGrids`:\n",
    "\n",
    "* [`spatial_gradient()`](phi/field/#phi.field.spatial_gradient) with `type=StaggeredGrid`\n",
    "* [`stagger()`](phi/field/#phi.field.stagger)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Values Tensor\n",
    "For non-periodic staggered grids, the `values` tensor is [non-uniform](https://tum-pbs.github.io/PhiFlow/Math.html#non-uniform-tensors)\n",
    "to reflect the different number of sample points for each component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[92m(xˢ=\u001b[94m(x=9, y=10)\u001b[92m, yˢ=\u001b[94m(x=10, y=9)\u001b[92m, vectorᶜ=x,y)\u001b[0m \u001b[94mconst 0.8027918338775635\u001b[0m"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grid.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Functions to get a uniform tensor:\n",
    "\n",
    "* [`uniform_values()`](phi/field/#phi.field.StaggeredGrid.uniform_values) interpolates the staggered values to the cell centers and returns a `CenteredGrid`\n",
    "* [`at_centers()`](phi/field/#phi.field.StaggeredGrid.at_centers) interpolates the staggered values to the cell centers and returns a `CenteredGrid`\n",
    "* [`staggered_tensor()`](phi/field/#phi.field.StaggeredGrid.staggered_tensor) pads the internal tensor to an invariant shape with n+1 entries along all dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[92m(xˢ=11, yˢ=11, vectorᶜ=x,y)\u001b[0m \u001b[94m0.597 ± 0.466\u001b[0m \u001b[37m(0e+00...1e+00)\u001b[0m"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grid.uniform_values()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Slicing\n",
    "Like tensors, grids can be sliced using the standard syntax.\n",
    "When selecting a vector component, such as `x` or `y`, the result is represented as a `CenteredGrid` with shifted locations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "grid.vector['x']  # select component"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Grids do not support slicing along spatial dimensions because the result would be ambiguous with StaggeredGrids.\n",
    "Instead, slice the `values` directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "grid.values.x[3:4]  # spatial slice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "grid.values.x[0]  # spatial slice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Slicing along batch dimensions has no special effect, this just slices the `values`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "grid.batch[0]  # batch slice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Fields can also be sliced using [`unstack()`](phi/field/#phi.field.unstack).\n",
    "This returns a `tuple` of all slices along a dimension."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
